from mpi4py import MPI
import numpy as np
from petsc4py import PETSc
from dolfinx import mesh, fem
from dolfinx.mesh import locate_entities_boundary, meshtags
from dolfinx.fem.petsc import assemble_matrix, assemble_vector
import ufl

# ----------------- 基本设置 -----------------
left_bound, right_bound = 0.0, 1.0
left_bound_value, right_bound_value = 1.0, 0.0

type_ = "SIPG"     # 本代码按 SIPG 正确号法实现；若切 NIPG/IIPG，需要相应改号
k = 2              # 多项式阶数
# 罚系数（保留你原来的变量名）
if type_ == "NIPG":
    epsilon = 1
    sigma_0 = 1
    sigma_1 = 0
elif type_ == "IIPG":
    epsilon = 0
    sigma_0 = 1
    sigma_1 = 1
elif type_ == "SIPG":
    epsilon = -1   # 注意：下面我们不再用 epsilon“翻号”，直接写 SIPG 的“负负正”
    sigma_0 = 20   # 增加罚参数以提高稳定性，对于k=2，建议至少10*k*(k+1)/2 ≈30，但取20测试
    sigma_1 = 0

h_list = [1/2, 1/4, 1/8, 1/16, 1/32]  # 匹配你的结果，添加1/2以重现nan问题

# 真解与右端（K=1）
def p(x):
    return (1 - x[0]) * np.exp(-x[0]**2)

def p_boundary(x):
    # 只在边界有值
    pass

def f(x):
    # f(x) = (4x^3 - 4x^2 - 6x + 2) e^{-x^2}
    return (4 * x[0]**3 - 4 * x[0]**2 - 6 * x[0] + 2) * np.exp(-x[0]**2)

# 误差收敛统计
beta1_list, beta2_list = [], []
beta1_err,  beta2_err  = [], []

# ----------------- 主循环 -----------------
for h in h_list:
    N = int(1/h)
    domain = mesh.create_interval(MPI.COMM_WORLD, N, [left_bound, right_bound])
    V = fem.functionspace(domain, ("DG", k))

    # 【重要】不要用 x 覆盖 PETSc 向量名，坐标变量改名为 X
    X = ufl.SpatialCoordinate(domain)

    u = ufl.TrialFunction(V)
    # 现在u + P1 = u_h  对应课本的P^DG_0  拆分了边界条件
    v = ufl.TestFunction(V)

    # -------- 测度与几何量 --------
    n   = ufl.FacetNormal(domain)
    dx  = ufl.dx(domain)
    dS  = ufl.dS(domain=domain)  # 内部面
    # 标记左右边界：1=left, 2=right（保留你的写法）
    fdim = domain.topology.dim - 1
    left_facets  = locate_entities_boundary(domain, fdim, lambda X: np.isclose(X[0], left_bound))
    right_facets = locate_entities_boundary(domain, fdim, lambda X: np.isclose(X[0], right_bound))
    tags = np.hstack([left_facets, right_facets]).astype(np.int32)
    vals = np.hstack([np.full_like(left_facets, 1, dtype=np.int32),
                      np.full_like(right_facets, 2, dtype=np.int32)])
    if tags.size == 0:
        facet_tags = meshtags(domain, fdim, np.array([], dtype=np.int32), np.array([], dtype=np.int32))
    else:
        order = np.argsort(tags)
        facet_tags = meshtags(domain, fdim, tags[order], vals[order])
    ds = ufl.Measure("ds", domain=domain, subdomain_data=facet_tags)
    ds_boundary = ds(1) + ds(2)

    # -------- 网格尺度与罚项尺度 --------
    # 【修正】内部面罚项用 h_avg；边界罚项用 h_cell
    h_cell = ufl.CellDiameter(domain)
    h_avg  = ufl.avg(h_cell)

    # -------- 右端插值 --------
    fh = fem.Function(V)
    fh.interpolate(lambda X: f(X))  # X 形状：(gdim, num_points)

    # -------- 定义 P1（连续 piecewise polynomial，degree k，但这里用 linear lift 1-x） --------
    V_cg = fem.functionspace(domain, ("Lagrange", k))  # CG space
    p1 = fem.Function(V_cg)
    p1.interpolate(lambda x: p_boundary(x)))

    # -------- DG 双线性型与右端 --------
    # 体项（K=1）
    a = ufl.inner(ufl.grad(u), ufl.grad(v)) * dx
    L = ufl.inner(fh, v) * dx - ufl.inner(ufl.grad(p1), ufl.grad(v)) * dx

    # ===== 内部面：SIPG 必须是 “负、负、正” =====
    a += (
        - ufl.inner(ufl.avg(ufl.grad(u)), ufl.jump(v, n)) * dS
        - ufl.inner(ufl.avg(ufl.grad(v)), ufl.jump(u, n)) * dS
        + (sigma_0 / h_avg) * ufl.inner(ufl.jump(u, n), ufl.jump(v, n)) * dS
    )

    # ===== 边界（强齐次 Dirichlet：u = 0）=====

    # -------- 组装与求解 --------
    A = assemble_matrix(fem.form(a))   # DG 不传 bcs
    A.assemble()
    b = assemble_vector(fem.form(L))
    b.assemble()

    uh = fem.Function(V)
    ksp = PETSc.KSP().create(domain.comm)
    ksp.setOperators(A)
    ksp.setType("preonly")
    ksp.getPC().setType("lu")

    x_sol = A.createVecRight()
    ksp.solve(b, x_sol)

    # 【修正】把解写回 uh；否则 uh 仍是零函数，误差全错
    uh.x.array[:] = x_sol.array

    # -------- 计算误差 --------
    # 精确解用几何坐标 X（不要用 x_sol 覆盖）
    p_exact = (1 - X[0]) * ufl.exp(-X[0]**2)
    # 【重要修正】全解为 uh (P0) + p1，全误差 e = uh + p1 - p_exact
    # 这个地方很麻烦 要让p1只在边界有值
    e = uh + p1 - p_exact

    # L2 误差
    L2_sq_loc = fem.assemble_scalar(fem.form(ufl.inner(e, e) * dx))
    L2_sq = domain.comm.allreduce(L2_sq_loc, op=MPI.SUM)
    L2_err = float(np.sqrt(L2_sq))

    # DG 能量范数：
    # 【修正】应包含三部分：体内梯度 + 内部跳量罚 + 边界罚（原代码漏了边界罚，但现在添加）
    En1 = fem.assemble_scalar(fem.form(ufl.inner(ufl.grad(e), ufl.grad(e)) * dx))
    En2 = fem.assemble_scalar(fem.form((sigma_0 / h_avg)  * ufl.inner(ufl.jump(e, n), ufl.jump(e, n)) * dS))
    En3 = fem.assemble_scalar(fem.form((sigma_0 / h_cell) * ufl.inner(e, e) * ds_boundary))
    En1 = domain.comm.allreduce(En1, op=MPI.SUM)
    En2 = domain.comm.allreduce(En2, op=MPI.SUM)
    En3 = domain.comm.allreduce(En3, op=MPI.SUM)
    print(f"En3 = {En3}.")
    E_err = float(np.sqrt(En1 + En2 + En3))

    if domain.comm.rank == 0:
        print(f"h={h: .5f} | L2 err={L2_err:.6e} | Energy err={E_err:.6e}")

    beta1_err.append(E_err)
    beta2_err.append(L2_err)

# -------- 收敛阶统计 --------
def _rates(errs, hs):
    out = []
    for i in range(len(errs) - 1):
        out.append(np.log(errs[i] / errs[i+1]) / np.log(hs[i] / hs[i+1]))
    return out

beta1_list = _rates(beta1_err, h_list)
beta2_list = _rates(beta2_err, h_list)

if MPI.COMM_WORLD.rank == 0:
    print("\n===== 收敛结果 =====")
    print("h:           ", ["{:.5f}".format(h) for h in h_list])
    print("能量范数 误差:", ["{:.3e}".format(e) for e in beta1_err])
    print("L2范数误差:", ["{:.3e}".format(e) for e in beta2_err])
    print("能量 (beta1):", ["{:.3f}".format(r) for r in beta1_list])
    print("L2 (beta2):", ["{:.3f}".format(r) for r in beta2_list])